import numpy as np

class AnnealingModel(object):
    
    @classmethod
    def get_age_r(cls, t, T, cproj):
        """t,T as output by Pecube"""
        dt = -np.diff(t)
        cT =  T[:-1] + np.diff(T)/2
        #cT =  T[1:]
        r = cls.rlen(dt, cT)
        a = cls.age(dt, r, cproj)
        return a, r

    @classmethod
    def rdens(cls, rlen,cproj):
        """Density reduction"""
        if cproj: # Ketcham, 2000, 2005
            dr = 1.6 * rlen - 0.6
            dr[rlen<0.765] = 9.205 * rlen[rlen<0.765]**2 - 9.157 * rlen[rlen<0.765] + 2.269
#        else: # Willett, 1997
#            dr = 2.862 * rlen - 1.2104
#            dr[rlen > 0.65 ] = rlen[rlen > 0.65]
#            dr[rlen < 0.423] = 0
        else: # MadTrax
            dr = 2.15 * rlen - 0.76
            dr[rlen > 0.65] = rlen[rlen > 0.65]
            dr[rlen < 0.35] = 0
        return dr


    @classmethod
    def age(cls, dt, rl, cproj):
        return np.sum(dt * cls.rdens(rl,cproj)) * 0.893 # Dens. reduction in standard (Willett 1992?)


class Laslett(AnnealingModel):
    """Laslett's FA model"""
    
    @classmethod
    def rlen(cls, dt, T):
        """Computes length reduction l/l0
        given len(Ts) dt-long steps from past to present.
        T[0] -- temp. before t[0]"""
        a  =  0.35
        b  =  2.7
        c0 = -4.87
        c1 =  0.000168
        c2 = 28.12
        c3 = 0
        t_eq = lambda r, T: np.exp(((((1 - r ** b) / b) ** a - 1) / a - c0) / c1 / T - c2)
        r_fa = lambda t, T: (1 - b * (a * (c0 + c1 * T * (np.log(t) + c2)) + 1) ** (1 / a)) ** (1 / b) # time in seconds
        
        #From MadTrax
        #c2 =  0.00472416
        #expos= lambda x, y: np.exp(y*np.log(x))
        #g    = lambda r: (expos((1.-expos(r,b))/b,a)-1.)/a
        #t_eq = lambda r, T: np.exp((-c2/c1)+((g(r)-c0)/c1)*(1./T-c3))
        #gr   = lambda t, T: c0+((c1*np.log(t)+c2)/((1./T)-c3)) 
        #r_fa = lambda t, T: expos(1.-b*expos(a*gr(t,T)+1.,1./a),1./b)
        
        
        
        dts = dt * 31556926000000
        teq = 0
        r = np.zeros(len(dt))
        TK = T + 273.15
        #TK = T[:-1] + np.diff(T)/2  + 273.15
        for i in reversed(range(len(TK))):
            r[i] = r_fa(dts[i] + teq, TK[i])
            if np.isnan(r[i]): r[i] = 0
            if r[i] < 0.13: break  #observ. bias Ketcham2000, p.6
            if i>0: teq = t_eq(r[i], TK[i-1])
        return r
    
    




class Ketcham2007FC(AnnealingModel):
    """Ketcham's FC 2007 model"""
    
    @classmethod
    def rlen(cls, dt, T):
        """Computes length reduction l/l0
        given len(Ts) dt-long steps from past to present.
        T[0] -- temp. before t[0]"""
        a  =  0.04672
        c0 =  0.39528
        c1 =  0.01073
        c2 =-65.12969
        c3 = -7.91715
        
        
        g = lambda r: (1/r - 1)**a
        ginv = lambda g: 1 / (g**(1/a) + 1)
        gr = lambda t, T: c0 + c1 * (np.log(t) - c2) / (np.log(1/T) - c3)
        
        t_eq = lambda gr, T: np.exp(c2 + (np.log(1/T) - c3) * (gr - c0) / c1)


        dts = dt * 31556926000000
        teq = 0
        r = np.zeros(len(dt))
        TK = T + 273.15
        #TK = T[:-1] + np.diff(T)/2  + 273.15
        for i in reversed(range(len(TK))):
            gr_i = gr(dts[i] + teq, TK[i])
            r[i] = ginv( gr_i )
            if np.isnan(r[i]): r[i] = 0
            if r[i] < 0.13: break  #observ. bias Ketcham2000, p.6
            if i>0: teq = t_eq(gr_i, TK[i-1])
        return r
    


